#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
证书生成脚本
用于生成双向TLS通信所需的CA证书、服务器证书和客户端证书
"""

import os
import subprocess
from pathlib import Path

def run_openssl_command(cmd):
    """执行OpenSSL命令"""
    try:
        result = subprocess.run(cmd, shell=True, check=True, capture_output=True, text=True)
        print(f"✓ 执行成功: {cmd}")
        return True
    except subprocess.CalledProcessError as e:
        print(f"✗ 执行失败: {cmd}")
        print(f"错误信息: {e.stderr}")
        return False

def create_certificates():
    """创建所有必需的证书"""
    # 创建证书目录
    certs_dir = Path("certs")
    certs_dir.mkdir(exist_ok=True)
    
    print("开始生成证书...")
    
    # 1. 生成CA私钥
    print("\n1. 生成CA私钥...")
    cmd = "openssl genrsa -out certs/ca-key.pem 4096"
    if not run_openssl_command(cmd):
        return False
    
    # 2. 生成CA证书
    print("\n2. 生成CA证书...")
    cmd = '''openssl req -new -x509 -days 365 -key certs/ca-key.pem -out certs/ca-cert.pem -subj "/C=CN/ST=Beijing/L=Beijing/O=Test/OU=Test/CN=TestCA"'''
    if not run_openssl_command(cmd):
        return False
    
    # 3. 生成服务器私钥
    print("\n3. 生成服务器私钥...")
    cmd = "openssl genrsa -out certs/server-key.pem 4096"
    if not run_openssl_command(cmd):
        return False
    
    # 4. 生成服务器证书签名请求
    print("\n4. 生成服务器证书签名请求...")
    cmd = '''openssl req -new -key certs/server-key.pem -out certs/server.csr -subj "/C=CN/ST=Beijing/L=Beijing/O=Test/OU=Test/CN=localhost"'''
    if not run_openssl_command(cmd):
        return False
    
    # 5. 用CA签名生成服务器证书
    print("\n5. 生成服务器证书...")
    cmd = "openssl x509 -req -days 365 -in certs/server.csr -CA certs/ca-cert.pem -CAkey certs/ca-key.pem -CAcreateserial -out certs/server-cert.pem"
    if not run_openssl_command(cmd):
        return False
    
    # 6. 生成客户端私钥
    print("\n6. 生成客户端私钥...")
    cmd = "openssl genrsa -out certs/client-key.pem 4096"
    if not run_openssl_command(cmd):
        return False
    
    # 7. 生成客户端证书签名请求
    print("\n7. 生成客户端证书签名请求...")
    cmd = '''openssl req -new -key certs/client-key.pem -out certs/client.csr -subj "/C=CN/ST=Beijing/L=Beijing/O=Test/OU=Test/CN=client"'''
    if not run_openssl_command(cmd):
        return False
    
    # 8. 用CA签名生成客户端证书
    print("\n8. 生成客户端证书...")
    cmd = "openssl x509 -req -days 365 -in certs/client.csr -CA certs/ca-cert.pem -CAkey certs/ca-key.pem -CAcreateserial -out certs/client-cert.pem"
    if not run_openssl_command(cmd):
        return False
    
    # 清理临时文件
    print("\n9. 清理临时文件...")
    try:
        os.remove("certs/server.csr")
        os.remove("certs/client.csr")
        print("✓ 临时文件清理完成")
    except:
        print("! 临时文件清理失败，请手动删除")
    
    print("\n✓ 所有证书生成完成！")
    print("\n生成的文件:")
    print("  certs/ca-cert.pem      - CA证书")
    print("  certs/ca-key.pem       - CA私钥")
    print("  certs/server-cert.pem  - 服务器证书")
    print("  certs/server-key.pem   - 服务器私钥")
    print("  certs/client-cert.pem  - 客户端证书")
    print("  certs/client-key.pem   - 客户端私钥")
    
    return True

def check_openssl():
    """检查OpenSSL是否可用"""
    try:
        result = subprocess.run(["openssl", "version"], capture_output=True, text=True)
        print(f"OpenSSL版本: {result.stdout.strip()}")
        return True
    except FileNotFoundError:
        print("错误: 未找到OpenSSL，请确保已安装OpenSSL并添加到PATH环境变量")
        return False

if __name__ == "__main__":
    print("=== TLS证书生成工具 ===")
    
    if not check_openssl():
        exit(1)
    
    if create_certificates():
        print("\n证书生成成功！现在可以运行TLS服务器和客户端了。")
    else:
        print("\n证书生成失败，请检查错误信息。")
        exit(1)